package com.qifei.xmlg_backend.controller;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okhttp3.RequestBody;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@RestController
@RequestMapping("/stream")
public class AiController {

    @Value("${guiji.api-url}")
    private String apiUrl;

    @Value("${guiji.api-key}")
    private String apiKey;

    private final OkHttpClient client = new OkHttpClient();
    private final ObjectMapper objectMapper = new ObjectMapper();

    @GetMapping("/sse")
    public SseEmitter stream(
            @RequestParam String prompt,
            @RequestParam String model,
            @RequestParam("max_tokens") int maxTokens,
            @RequestParam double temperature,
            @RequestParam("top_p") double topP,
            @RequestParam("top_k") int topK,
            @RequestParam("frequency_penalty") double frequencyPenalty
    ) {
        SseEmitter emitter = new SseEmitter(0L);

        new Thread(() -> {
            try {
                // ✅ 用传参构造请求体
                Map<String, Object> requestBodyMap = new HashMap<>();
                requestBodyMap.put("model", model);
                requestBodyMap.put("stream", true);
                requestBodyMap.put("max_tokens", maxTokens);
                requestBodyMap.put("temperature", temperature);
                requestBodyMap.put("top_p", topP);
                requestBodyMap.put("top_k", topK);
                requestBodyMap.put("frequency_penalty", frequencyPenalty);
                requestBodyMap.put("n", 1);
                requestBodyMap.put("stop", List.of());
                requestBodyMap.put("messages", List.of(
                        Map.of("role", "user", "content", prompt)
                ));
                log.info(requestBodyMap.toString());
                String jsonBody = objectMapper.writeValueAsString(requestBodyMap);

                RequestBody body = RequestBody.create(
                        jsonBody,
                        MediaType.get("application/json")
                );

                Request request = new Request.Builder()
                        .url(apiUrl)
                        .header("Authorization", "Bearer " + apiKey)
                        .post(body)
                        .build();

                Response response = client.newCall(request).execute();

                if (!response.isSuccessful()) {
                    emitter.send(SseEmitter.event().data("Error: " + response));
                    emitter.complete();
                    return;
                }

                BufferedReader reader = new BufferedReader(new InputStreamReader(response.body().byteStream()));
                String line;
                while ((line = reader.readLine()) != null) {
                    if (line.startsWith("data: ")) {
                        String json = line.substring("data: ".length());
                        emitter.send(SseEmitter.event().data(json));
                        if (json.equals("[DONE]")) {
                            break;
                        }
                    }
                }

                emitter.complete();
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        }).start();

        return emitter;
    }
}
